#!/usr/bin/python3.2
from urllib.request import Request, urlopen
import os
import gzip
from urllib.parse import urlparse
import threading
import time
import sys
from queue import Queue
import datetime

def urljoin(*pieces):
    return '/'.join(s.strip('/') for s in pieces)



def wget(base, *parts):
    localdir = os.path.join(*parts[:-1])
    
    try:
        os.makedirs(localdir)
    except OSError as e:
        if e.errno != 17:
            raise
    
    localpath = os.path.join(*parts)
    addr = urljoin(base, *parts)
    resp = urlopen(addr)
    data = resp.read()
    
    with open(localpath,'wb') as f:
        f.write(data)
    
    return localpath


# http://manpages.ubuntu.com/manpages/hardy/man5/deb-version.5.html
class DebVersion(object):
    def __init__(self, epoch, upstream, debian):
        self.epoch = epoch
        self.upstream = upstream
        self.debian = debian

    @staticmethod
    def from_version_string(version_string):
        return DebVersion(*DebVersion.parse_version_string(version_string))

    @staticmethod
    def parse_version_string(version_string):
        epoch = None
        version = None
        debian = None
        
        if ':' in version_string:
            epoch = version_string.split(':')[0]
            version_string = version_string[len(epoch)+1:]
        
        if '-' in version_string:
            debian = version_string.split('-')[-1]
            version_string = version_string[:-(len(debian)+1)]
        
        version = version_string
        
        return (epoch, version, debian)

    def __str__(self):
        s = self.upstream
        
        if self.epoch is not None:
            s = '%s:%s' % (self.epoch, s)
        
        if self.debian is not None:
            s = '%s-%s' % (s, self.debian)

        return s
    
    def __repr__(self):
        return '%s(%s,%s,%s)' % (self.__class__.__name__, repr(self.epoch), repr(self.upstream), repr(self.debian))



class Dependency(object):
    def __init__(self, depends_entry):
        self.package = depends_entry.split(' ')[0]
    


class Package(object):
    def __init__(self, package_entry):
        
        package_entry = package_entry.split('\n')
        
        for line in package_entry:
            attrname = line.split(': ')[0].replace('-','_')
            attrdata = line[len(attrname+': '):]
            setattr(self, attrname, attrdata)
        
        if hasattr(self,'Depends'):
                deps = self.Depends
                self.Depends = []
                
                for dep in deps.split(', '):
                    self.Depends = Dependency(dep)
        
        if hasattr(self, 'Version'):
                self.Version = DebVersion.from_version_string(self.Version)



class ReleaseInfoFile(object):
    def __init__(self, hashtype, line):
        line = line.split(' ')
        
        while '' in line:
            line.remove('')
        
        self.hashtype = hashtype
        self.hash = line[0]
        self.size = line[1]
        self.path = line[2]



class ReleaseInfo(object):
    def __init__(self, release):
        data = open(release,'r').read().split('\n')
        
        i = 0
        
        while i < len(data):
            if data[i].split(':')[0] in ['Origin', 'Label', 'Suite', 'Version', 'Codename', 'Date', 'Description']:
                attrname = data[i].split(': ')[0]
                attrdata = data[i][len(attrname)+len(': '):]
                setattr(self, attrname, attrdata)
                i += 1
            
            elif data[i] in ['MD5Sum:', 'SHA1:', 'SHA256:']:
                attrname = data[i].rstrip(':')
                attrdata = []
                i += 1
                
                while i < len(data) and len(data[i]) > 0 and data[i][0] == ' ':
                    attrdata.append(ReleaseInfoFile(attrname, data[i]))
                    i += 1
                
                setattr(self, attrname, attrdata)
            
            elif data[i].split(':')[0] in ['Architectures','Components']:
                setattr(self, data[i].split(':')[0], data[i].split(': ')[1].split(' '))
                i += 1
            
            else:
                i += 1
    
    def contains_all_for_component(self, component):
        search_path = urljoin(component,'binary-all','Packages.gz')
        
        for entry in self.MD5Sum:
            if entry.path == search_path:
                return True
        
        return False



def get_package_list(repo,dist,component,arch):
    return wget(repo, 'dists', dist, component, 'binary-'+arch, 'Packages.gz')



def get_gpg_signature(repo, dist):
    return wget(repo,'dists',dist,'Release.gpg')
    


def read_package_list(localpath):
    f = gzip.open(localpath, 'r')
    data = str(f.read(),'utf-8')
    data = data.split('\n\n')
    
    while '' in data:
        data.remove('')
    
    for i in range(len(data)):
        data[i] = Package(data[i])
    
    return data



def top_release_path(repo, dist):
    return urljoin(repo,'dists',dist,'Release')



def get_release_info(repo, dist):
    return wget(repo,'dists',dist,'Release')



if __name__ == '__main__':
    from argparse import ArgumentParser

    parser = ArgumentParser()
    parser.add_argument('-r', '--repository', action='store', help='The address to the repository.', default='http://http.debian.net/debian')
    parser.add_argument('-d', '--distribution', action='store', help='The distribution to use.', default='wheezy')
    parser.add_argument('-c', '--component', action='store', help='The component to download. Exclude this option to grab all components.', default=None, nargs='+')
    parser.add_argument('-a', '--architecture', action='store', help='CPU Architecture.', default=None, nargs='+')
    parser.add_argument('-t', '--threads', action='store', help='Use multiple threads when downloading the packages.', type=int, default=1)
    parser.add_argument('-q', '--quit', action='store_true', help='Quit before downloading the packages.', default=False)
    options = parser.parse_args()
    
    print('Fetching GPG signature...')
    get_gpg_signature(options.repository, options.distribution)
    
    
    print('Getting release information...')
    release_info_file = get_release_info(options.repository, options.distribution)
    
    
    print('Parsing release information...')
    release_info = ReleaseInfo(release_info_file)
    
    
    if options.architecture is not None:
        exitcode = 0
        
        for arch in options.architecture:
            if arch not in release_info.Architectures:
                print('Bad architecture, "%s" is not a valid architecture. The repo has support for:\n    %s' % (options.architecture, ', '.join(release_info.Architectures)))
                exitcode = -1
        
        if exitcode:
            exit(exitcode)
    
    elif options.architecture is None:
        options.architecture = release_info.Architectures
    
    
    if options.component is not None:
        exitcode = 0
        
        for comp in options.component:
            if comp not in release_info.Components:
                print('Bad component, "%s" is not a valid component. The repo has support for:\n    %s' % (options.component, ', '.join(release_info.Components)))
                exitcode = -1
        
        if exitcode:
            exit(exitcode)

    elif options.component is None:
        options.component = release_info.Components
    
    
    packages_to_download = set()
    
    for component in options.component:
        architectures = list(options.architecture)
        
        if release_info.contains_all_for_component(component):
            architectures.append('all')
        
        for architecture in architectures:
            print('Downloading package list for %s/%s' % (component,architecture))
            packages_gz = get_package_list(options.repository,options.distribution, component, architecture)
            
            
            print('Parsing package list...')
            packages = read_package_list(packages_gz)
            
            for package in packages:
                packages_to_download.add(package)
    
    total_size = 0
    for pkg in packages_to_download:
        if hasattr(pkg, 'Size'):
            total_size += int(pkg.Size)
    
    if total_size > 1024*1024*1024:
        print_size = '%s GB' % (round(total_size / (1024*1024*1024),2))
    
    elif total_size > 1024*1024:
        print_size = '%s MB' % (round(total_size / (1024*1024), 2))
    
    elif total_size > 1024:
        print_size = '%s KB' % (round(total_size / 1024,2))
    
    print('\nPackages Found:\n    Count: %d\n    Total Size: %s\n' % (len(packages_to_download), print_size))
    
    if options.quit:
        print('Exiting...')
        exit(0)
    
    print('Starting download...')
    
    class CallbackThread(threading.Thread):
        def __init__(self, *args, **kwargs):
            threading.Thread.__init__(self, *args, **kwargs)
            self.callback = kwargs.get('callback', None)
        
        def run(self, *args, **kwargs):
            self._target(*self._args, **self._kwargs)
            if self.callback is not None:
                self.callback()
            return
    
    def create_closures(max_thread):
        open_threads = [0]
        
        def inc():
            open_threads[0] += 1
        
        def dec():
            open_threads[0] -= 1
        
        def check():
            return open_threads[0]
        
        return inc, dec, check
    
    def get_speed(start_time, bytes_downloaded):
        
        delta_time = datetime.datetime.now() - start_time
        if delta_time.seconds == 0:
            return '0 Bps'
        
        download_speed = total_size / delta_time.seconds
    
        if download_speed > 1024*1024*1024:
            print_speed = '%s GBps' % (round(download_speed / (1024*1024*1024),2))
        
        elif download_speed > 1024*1024:
            print_speed = '%s MBps' % (round(download_speed / (1024*1024), 2))
        
        elif download_speed > 1024:
            print_speed = '%s KBps' % (round(download_speed / 1024,2))
        
        else:
            print_speed = '%s Bps' % (download_speed)
        
        return print_speed
    
    def joiner(thread_count, q, start_time):
        threads_joined = 0
        data_downloaded = 0
        
        sys.stdout.write('\r%d/%d (%s)' % (threads_joined, thread_count, get_speed(start_time, data_downloaded) ))
        
        while threads_joined < thread_count:
            t = q.get()
            t.join()
            dec()
            threads_joined += 1
            data_downloaded += int(t.pkg.Size)
            sys.stdout.write('\r%d/%d (%s)' % (threads_joined, thread_count, get_speed(start_time, data_downloaded) ))
    
    
    inc, dec, check = create_closures(options.threads)
    join_queue = Queue()
    
    
    start_time = datetime.datetime.now()
    
    # Start a thread to cleanup the other threads as they complete.
    joiner_thread = threading.Thread(target=joiner, args=(len(packages_to_download), join_queue, start_time))
    joiner_thread.daemon = False
    joiner_thread.start()
    
    for pkg in packages_to_download:
        addr = urljoin(options.repository, pkg.Filename)
        
        while check() >= options.threads:
            time.sleep(0.1)
        
        inc()
        t = CallbackThread(target=wget, args=[options.repository] + pkg.Filename.split('/'))
        t.callback = lambda t=t: join_queue.put(t) # Have the thread queue itself for joining, when its done.
        t.daemon = False
        t.pkg = pkg
        t.start()
    
    joiner_thread.join()
    
    total_time = datetime.datetime.now() - start_time
    print('')
    print('Downloaded %s in %d seconds. Average speed was %dMBps.' % (print_size, total_time.seconds, (total_size/(1024*1024))/total_time.seconds))
























